<!DOCTYPE html>
<html lang="en">
  <head>
    <meta charset="UTF-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>手写数字计算器</title>
    <style>
      body {
        font-family: "Arial", sans-serif;
        text-align: center;
        max-width: 1200px;
        margin: 0 auto;
        padding: 20px;
        background-color: #f5f5f5;
      }

      h1 {
        color: #333;
        margin-bottom: 25px;
      }

      .container {
        display: flex;
        justify-content: space-between;
        flex-wrap: wrap;
      }

      .canvas-container {
        width: 45%;
        margin-bottom: 20px;
      }

      .control-panel {
        width: 100%;
        margin: 20px 0;
        padding: 20px;
        background-color: white;
        border-radius: 8px;
        box-shadow: 0 2px 10px rgba(0, 0, 0, 0.1);
      }

      button {
        background-color: #4caf50;
        border: none;
        color: white;
        padding: 10px 20px;
        text-align: center;
        text-decoration: none;
        display: inline-block;
        font-size: 16px;
        margin: 10px 5px;
        cursor: pointer;
        border-radius: 4px;
        transition: background-color 0.3s;
      }

      button:hover {
        background-color: #45a049;
      }

      .result-container {
        margin-top: 20px;
        font-size: 24px;
        font-weight: bold;
        background-color: #e0f7fa;
        padding: 15px;
        border-radius: 8px;
      }

      canvas {
        border: 2px solid #333;
        background-color: white;
        cursor: pointer;
        border-radius: 8px;
        box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
      }

      .digit-preview {
        margin-top: 10px;
        font-weight: bold;
      }
    </style>
  </head>
  <body>
    <h1>手写数字计算器</h1>

    <p>在画布上手写数字，然后点击"计算和"按钮查看结果</p>

    <div class="control-panel">
      <button id="clearAll">清除所有</button>
    </div>
    <div id="sampleContainer">
      <h3>样本采集</h3>
      <div class="canvas-container" style="width: 100%">
        <canvas id="digitCanvas3" width="200" height="200"></canvas>
        <canvas id="digitCanvas4" width="200" height="200"></canvas>
        <canvas id="digitCanvas5" width="200" height="200"></canvas>
        <canvas id="digitCanvas6" width="200" height="200"></canvas>
        <canvas id="digitCanvas7" width="200" height="200"></canvas>
        <canvas id="digitCanvas8" width="200" height="200"></canvas>
        <canvas id="digitCanvas9" width="200" height="200"></canvas>
        <canvas id="digitCanvas10" width="200" height="200"></canvas>
        <canvas id="digitCanvas11" width="200" height="200"></canvas>
        <canvas id="digitCanvas12" width="200" height="200"></canvas>
      </div>
      <div style="display: flex; align-items: center; justify-content: center">
        <button id="collectSamples">采集样本</button>
        <button id="clearSamples">清除样本</button>
      </div>
    </div>
    <div>
      <h3>重新训练模型</h3>
      <input type="file" id="sampleFile" accept=".json" multiple />
    </div>
    <div>
      <h3>使用新模型</h3>
      <input type="file" id="uploadModel" />
    </div>
    <div class="container">
      <div class="canvas-container">
        <h3>第一个数字</h3>
        <canvas id="digitCanvas1" width="200" height="200"></canvas>
        <div class="digit-preview" id="digitResult1">识别结果：-</div>
      </div>
      <div class="canvas-container">
        <h3>第二个数字</h3>
        <canvas id="digitCanvas2" width="200" height="200"></canvas>
        <div class="digit-preview" id="digitResult2">识别结果：-</div>
      </div>
      <div style="display: flex; align-items: center; justify-content: center; flex-direction: column">
        <button id="clearInputs">清除数字</button>
         <button id="calculate">计算和</button>
      </div>
    </div>
    <div class="result-container">
      <p id="resultDisplay">等待计算...</p>
    </div>

    <script src="././cnn_simple_hwr/tfjs-vis.umd.min.js"></script>
    <script src="././cnn_simple_hwr/tf.min.js"></script>

    <script type="module">
      // 模型地址
      const MODEL_URL = "./my_model/model.json"; // 使用相对路径指向模型文件

      let model;
      let isDrawing = false;
      let lastPoint = {};
      const canvases = {
        1: document.getElementById("digitCanvas1"),
        2: document.getElementById("digitCanvas2"),
        3: document.getElementById("digitCanvas3"),
        4: document.getElementById("digitCanvas4"),
        5: document.getElementById("digitCanvas5"),
        6: document.getElementById("digitCanvas6"),
        7: document.getElementById("digitCanvas7"),
        8: document.getElementById("digitCanvas8"),
        9: document.getElementById("digitCanvas9"),
        10: document.getElementById("digitCanvas10"),
        11: document.getElementById("digitCanvas11"),
        12: document.getElementById("digitCanvas12"),
      };
      const contexts = {
        1: canvases[1].getContext("2d"),
        2: canvases[2].getContext("2d"),
        3: canvases[3].getContext("2d"),
        4: canvases[4].getContext("2d"),
        5: canvases[5].getContext("2d"),
        6: canvases[6].getContext("2d"),
        7: canvases[7].getContext("2d"),
        8: canvases[8].getContext("2d"),
        9: canvases[9].getContext("2d"),
        10: canvases[10].getContext("2d"),
        11: canvases[11].getContext("2d"),
        12: canvases[12].getContext("2d"),
      };

      // 初始化画布
      function initCanvas() {
        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12].forEach((index) => {
          contexts[index].fillStyle = "white";
          contexts[index].fillRect(
            0,
            0,
            canvases[index].width,
            canvases[index].height
          );
        });

        setupDrawingEvents();
      }

      // 设置画布事件监听
      function setupDrawingEvents() {
        // 为每个画布添加事件监听
        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12].forEach((index) => {
          const canvas = canvases[index];
          const context = contexts[index];

          // 开始绘画
          canvas.addEventListener("mousedown", (e) => {
            isDrawing = true;
            const rect = canvas.getBoundingClientRect();
            lastPoint[index] = {
              x: e.clientX - rect.left,
              y: e.clientY - rect.top,
            };
            context.beginPath();
            context.moveTo(lastPoint[index].x, lastPoint[index].y);
          });

          canvas.addEventListener("touchstart", (e) => {
            isDrawing = true;
            const rect = canvas.getBoundingClientRect();
            lastPoint[index] = {
              x: e.touches[0].clientX - rect.left,
              y: e.touches[0].clientY - rect.top,
            };
            context.beginPath();
            context.moveTo(lastPoint[index].x, lastPoint[index].y);
            e.preventDefault();
          });

          // 绘画中
          canvas.addEventListener("mousemove", (e) => {
            if (isDrawing) {
              draw(e, context, canvas);
            }
          });

          canvas.addEventListener("touchmove", (e) => {
            if (isDrawing) {
              draw(e, context, canvas);
            }
            e.preventDefault();
          });

          // 结束绘画
          canvas.addEventListener("mouseup", () => {
            isDrawing = false;
            if (index >= 3) return;
            recognizeDigit(index);
          });

          canvas.addEventListener("touchend", () => {
            isDrawing = false;
            if (index >= 3) return;
            recognizeDigit(index);
            event.preventDefault();
          });

          canvas.addEventListener("mouseleave", () => {
            isDrawing = false;
          });
          canvas.addEventListener("touchcancel", () => {
            isDrawing = false;
          });
        });
      }

      // 绘制函数
      // In draw function:
      function draw(event, context, canvas) {
        let rect = canvas.getBoundingClientRect();
        let x, y;

        if (event.type.includes("touch")) {
          x = event.touches[0].clientX - rect.left;
          y = event.touches[0].clientY - rect.top;
        } else {
          x = event.clientX - rect.left;
          y = event.clientY - rect.top;
        }

        context.lineWidth = 15;
        context.lineCap = "round";
        context.strokeStyle = "black";

        context.lineTo(x, y);
        context.stroke();
        context.beginPath();
        context.moveTo(x, y);
      }

      // 清除指定画布
      function clearCanvas(canvasId) {
        const context = contexts[canvasId];
        context.fillStyle = "white";
        context.fillRect(
          0,
          0,
          canvases[canvasId].width,
          canvases[canvasId].height
        );
        const resultEl = document.getElementById(`digitResult${canvasId}`);
        if (canvasId < 3) {
          resultEl.textContent = "识别结果：-";
        }
      }

      // 清除所有画布
      function clearAllCanvases() {
        clearCanvas(1);
        clearCanvas(2);
        clearCanvas(3);
        clearCanvas(4);
        clearCanvas(5);
        clearCanvas(6);
        clearCanvas(7);
        clearCanvas(8);
        clearCanvas(9);
        clearCanvas(10);
        clearCanvas(11);
        clearCanvas(12);

        document.getElementById("resultDisplay").textContent = "等待计算...";
      }

      // 加载模型
      async function loadModel() {
        console.log("正在加载模型...");
        model = await tf.loadLayersModel(MODEL_URL);
        console.log("模型加载完成");
      }

      // 识别数字
      async function recognizeDigit(canvasId) {
        if (!model) return;

        const canvas = canvases[canvasId];
        const digitElement = document.getElementById(`digitResult${canvasId}`);

        // 复制图像到新画布进行预处理
        const tempCanvas = document.createElement("canvas");
        const tempCtx = tempCanvas.getContext("2d");
        tempCanvas.width = 28;
        tempCanvas.height = 28;

        // 绘制到临时画布并调整大小
        tempCtx.drawImage(canvas, 0, 0, 28, 28);

        // 获取像素数据并归一化
        const imgData = tempCtx.getImageData(0, 0, 28, 28);
        const input = tf.tidy(() => {
          let img = tf.browser.fromPixels(imgData, 1);
          img = tf.abs(tf.sub(1, img.toFloat().div(255.0))); // Invert: black digit on white to white digit on black
          img = img.reshape([1, 28, 28, 1]);
          return img;
        });

        const prediction = model.predict(input);
        const digit = prediction.argMax(1).dataSync()[0];

        digitElement.textContent = `识别结果：${digit}`;
        input.dispose();
        prediction.dispose();
        return digit;
      }

      // 计算两个数字的和
      async function calculateSum() {
        if (!model) {
          alert("模型尚未加载完成，请稍候...");
          return;
        }

        const digit1 = await recognizeDigit(1);
        const digit2 = await recognizeDigit(2);

        if (digit1 === undefined || digit2 === undefined) {
          alert("请先在两个画布上书写数字");
          return;
        }

        const sum = digit1 + digit2;
        document.getElementById(
          "resultDisplay"
        ).textContent = `计算结果：${digit1} + ${digit2} = ${sum}`;
      }

      // 采集样本
      const sampleCounters = {};
      function saveCanvasSample(canvasId, label) {
        if (!sampleCounters[label]) {
          sampleCounters[label] = 1;
        } else {
          sampleCounters[label] += 1;
        }

        const canvas = canvases[canvasId];
        const tempCanvas = document.createElement("canvas");
        tempCanvas.width = 28;
        tempCanvas.height = 28;
        const tempCtx = tempCanvas.getContext("2d");
        tempCtx.drawImage(canvas, 0, 0, 28, 28);

        // Get image data and preprocess (invert, normalize)
        const imgData = tempCtx.getImageData(0, 0, 28, 28);
        let img = tf.browser.fromPixels(imgData, 1);
        img = tf.abs(tf.sub(1, img.toFloat().div(255.0))); // Invert if needed
        img = img.reshape([28, 28]);
        const arr = Array.from(img.dataSync());

        // Save as JSON (image array + label)
        const sample = {
          image: arr,
          label: label,
        };
        const blob = new Blob([JSON.stringify(sample)], {
          type: "application/json",
        });
        const url = URL.createObjectURL(blob);

        // Download the sample
        const a = document.createElement("a");
        a.href = url;
        a.download = `digit_sample_${label}_${sampleCounters[label]}.json`;

        a.click();
        URL.revokeObjectURL(url);

        img.dispose();
      }

      async function loadSamples(files) {
        const images = [];
        const labels = [];
        for (const file of files) {
          const fileName = file.name;
          const text = await file.text();
          const sample = JSON.parse(text);
          images.push(sample.image);
          labels.push(sample.label);
        }
        // Convert to tensors
        const xs = tf
          .tensor2d(images, [images.length, 28 * 28])
          .reshape([images.length, 28, 28, 1]);
        const ys = tf.oneHot(tf.tensor1d(labels, "int32"), 10);
        return { xs, ys };
      }

      document
        .getElementById("sampleFile")
        .addEventListener("change", async (e) => {
          if (!model) {
            alert("模型尚未加载完成，请稍候...");
            return;
          }
          const files = e.target.files;
          const { xs, ys } = await loadSamples(files);

          console.log("加载的样本数量:", xs.shape[0]);
          console.log("加载的标签数量:", ys.shape[0]);

          // Compile the model with a new optimizer
          model.compile({
            optimizer: tf.train.adam(0.001, 0.9, 0.999, 1e-7),
            loss: "categoricalCrossentropy",
            metrics: ["accuracy"],
          });

          // Fine-tune the model
          await model.fit(xs, ys, {
            epochs: 10,
            batchSize: 16,
            shuffle: true,
            callbacks: tfvis.show.fitCallbacks(
              { name: "Fine-tuning", tab: "Training" },
              ["loss", "acc"],
              { height: 200 }
            ),
          });

          // Optionally save the updated model
          await model.save("downloads://fine_tuned_model");
        });

      // 上传新模型
      document
        .getElementById("uploadModel")
        .addEventListener("change", async (e) => {
          const file = e.target.files[0];
          if (!file) {
            alert("请选择一个模型文件！");
            return;
          }
          model = await tf.loadLayersModel(`./my_model/${file.name}`);

          alert("已加载新的模型，后续识别将使用新模型！");
        });
      
      // 事件绑定
      document
        .getElementById("clearAll")
        .addEventListener("click", clearAllCanvases);
      document
        .getElementById("calculate")
        .addEventListener("click", calculateSum);
      document.getElementById("clearSamples").addEventListener("click", () => {
        for (let i = 3; i <= 12; i++) {
          clearCanvas(i);
        }
        alert("所有样本已清除");
        for (let i = 3; i <= 12; i++) {
          sampleCounters[i] = 0;
        }
      });

      document
        .getElementById("collectSamples")
        .addEventListener("click", () => {
          const label = prompt("请输入你画的数字标签 (0-9):");
          if (label !== null && label.match(/^[0-9]$/)) {
            for (let i = 3; i <= 12; i++) {
              saveCanvasSample(i, parseInt(label));
            }
          } else {
            alert("请输入0-9之间的数字！");
          }
        });

      // 清除输入数字
      document.getElementById("clearInputs").addEventListener("click", () => {
        clearCanvas(1);
        clearCanvas(2);
        document.getElementById("digitResult1").textContent = "识别结果：-";
        document.getElementById("digitResult2").textContent = "识别结果：-";
        document.getElementById("resultDisplay").textContent = "等待计算...";
      });

      // 初始化应用
      window.onload = async function () {
        initCanvas();
        await loadModel();
      };
    </script>
  </body>
</html>
